// test_numa.c
#include "test_framework.h"
#include "../hisi_l0_mem_pool.h"
#include "test_common.h"
#include "test_registration.h"

/**
 * test_numa_node_allocation - Tests allocation on different NUMA nodes
 * 
 * This test verifies:
 * 1. NUMA-aware allocation functionality
 * 2. Allocation on different NUMA nodes
 * 3. System behavior with multiple NUMA nodes
 * 4. Strict validation of all allocation and deallocation operations
 * 
 * Tests NUMA node allocation capabilities with strict validation
 */
static int test_numa_node_allocation(struct seq_file *m)
{
    void *ptr;
    int num_nodes = num_online_nodes();
    int i, j;
    int passed = 0, total = 0;
    bool has_failures = false;
    bool content_valid;
    unsigned char *data;
    
    // Limit to reasonable number
    if (num_nodes > 4) num_nodes = 4;
    
    for (i = 0; i < num_nodes; i++) {
        total++;
        ptr = l0_kmalloc(512, i);
        if (ptr) {
            // Verify memory is accessible
            memset(ptr, i & 0xFF, 512);
            
            // Verify content
            data = (unsigned char *)ptr;
            content_valid = true;
            for (j = 0; j < 512; j++) {
                if (data[j] != (i & 0xFF)) {
                    content_valid = false;
                    break;
                }
            }
            
            if (content_valid) {
                l0_kfree(ptr);
                TEST_PASS_DEBUG(m, "NUMA node %d allocation and content verification successful", i);
                passed++;
            } else {
                l0_kfree(ptr);
                TEST_FAIL(m, "NUMA node %d allocation content verification failed", i);
                has_failures = true;
            }
        } else {
            TEST_FAIL(m, "NUMA node %d allocation failed", i);
            has_failures = true;
        }
    }
    
    TEST_END(m, "NUMA node allocation tests: %d/%d operations passed", passed, total);
    
    // For strict validation, we require all critical operations to pass
    if (has_failures) {
        TEST_FAIL(m, "Some operations failed during NUMA node allocation tests");
        return -1;
    }
    
    return (passed = total) ? 0 : -1;
}

/**
 * test_numa_node_migration - Tests memory accessibility across NUMA nodes
 * 
 * This test verifies:
 * 1. Memory allocated on one NUMA node is accessible from another
 * 2. NUMA memory coherency
 * 3. Cross-node memory operations
 * 4. Strict validation of all allocation and deallocation operations
 * 
 * Tests NUMA memory consistency and accessibility with strict validation
 */
static int test_numa_node_migration(struct seq_file *m)
{
    void *ptrs[20];
    int num_nodes = num_online_nodes();
    int i, j;
    int passed = 0, total = 0;
    bool has_failures = false;
    unsigned char *data;
    bool content_valid;
    unsigned char first_byte;
    unsigned char middle_byte;
    unsigned char last_byte;
    
    // Limit to reasonable number
    if (num_nodes > 4) num_nodes = 4;
    
    if (num_nodes < 2) {
        TEST_INFO(m, "NUMA migration test skipped: less than 2 nodes available");
        return 0;
    }
    
    // Allocate memory on different nodes and verify accessibility
    for (i = 0; i < num_nodes && i < ARRAY_SIZE(ptrs); i++) {
        total++;
        ptrs[i] = l0_kmalloc(512, i); // Allocate on node i
        if (ptrs[i]) {
            // Fill with node-specific pattern
            memset(ptrs[i], i & 0xFF, 512);
            TEST_PASS_DEBUG(m, "NUMA node %d allocation and content write successful", i);
            passed++;
        } else {
            TEST_FAIL(m, "NUMA node %d allocation failed", i);
            has_failures = true;
            // Initialize to NULL to avoid issues in cleanup
            ptrs[i] = NULL;
        }
    }
    
    // Verify all allocated memory is accessible and contains correct data
    for (i = 0; i < num_nodes && i < ARRAY_SIZE(ptrs); i++) {
        total++;
        if (ptrs[i]) {
            // Access memory and verify content
            data = (unsigned char *)ptrs[i];
            content_valid = true;
            
            // Check that all bytes have the correct pattern
            for (j = 0; j < 512; j++) {
                if (data[j] != (i & 0xFF)) {
                    content_valid = false;
                    break;
                }
            }
            
            if (content_valid) {
                TEST_PASS_DEBUG(m, "NUMA node %d memory content verification successful", i);
                passed++;
            } else {
                TEST_FAIL(m, "Data corruption detected in NUMA node %d memory", i);
                has_failures = true;
            }
        }
    }
    
    // Try to access all allocated memory from a different node context
    for (i = 0; i < num_nodes && i < ARRAY_SIZE(ptrs); i++) {
        total++;
        if (ptrs[i]) {
            // Access from test context node (potentially different from allocation node)
            data = (unsigned char *)ptrs[i];
            first_byte = data[0];
            middle_byte = data[256];
            last_byte = data[511];
            
            // Verify these bytes have the expected values
            if (first_byte == (i & 0xFF) && 
                middle_byte == (i & 0xFF) && 
                last_byte == (i & 0xFF)) {
                TEST_PASS_DEBUG(m, "Cross-node access to NUMA node %d memory successful", i);
                passed++;
            } else {
                TEST_FAIL(m, "Cross-node access to NUMA node %d memory failed", i);
                has_failures = true;
            }
        }
    }
    
    // Free all allocated memory
    for (i = 0; i < num_nodes && i < ARRAY_SIZE(ptrs); i++) {
        total++;
        if (ptrs[i]) {
            l0_kfree(ptrs[i]);
            TEST_PASS_DEBUG(m, "NUMA node %d memory deallocation successful", i);
            passed++;
        }
    }
    
    TEST_END(m, "NUMA node migration test:");
    TEST_END(m, "  Allocated and verified memory on %d nodes", num_nodes);
    TEST_END(m, "  Passed: %d/%d operations", passed, total);
    
    // For strict validation, we require all critical operations to pass
    if (has_failures) {
        TEST_FAIL(m, "Some operations failed during NUMA node migration tests");
        return -1;
    }
    
    if (passed > 0) {
        TEST_PASS(m, "NUMA node migration test completed successfully");
        return 0;
    } else {
        TEST_FAIL(m, "NUMA node migration test failed");
        return -1;
    }
}


/**
 * test_numa_locality - Tests NUMA locality of allocations
 * 
 * This test verifies:
 * 1. Memory is allocated from the requested NUMA node when possible
 * 2. Proper fallback when requested node is unavailable
 * 3. NUMA node assignment consistency
 * 
 * Tests NUMA locality guarantees
 */
static int test_numa_locality(struct seq_file *m)
{
    void *ptr;
    int num_nodes = num_online_nodes();
    int i;
    int passed = 0, total = 0;
    
    if (num_nodes < 2) {
        TEST_INFO(m, "NUMA locality test skipped: less than 2 nodes available");
        return 0;
    }
    
    // Limit to reasonable number
    if (num_nodes > 4) num_nodes = 4;
    
    for (i = 0; i < num_nodes; i++) {
        total++;
        ptr = l0_kmalloc(512, i);
        if (ptr) {
            // We can't directly verify NUMA node of allocated memory in kernel space
            // but we can verify the allocation succeeded
            memset(ptr, i & 0xFF, 512);
            l0_kfree(ptr);
            TEST_PASS_DEBUG(m, "NUMA node %d allocation successful", i);
            passed++;
        } else {
            TEST_FAIL(m, "NUMA node %d allocation failed", i);
        }
    }
    
    // Test with invalid node ID (should fallback)
    total++;
    ptr = l0_kmalloc(512, num_nodes + 1);
    if (ptr) {
        l0_kfree(ptr);
        TEST_FAIL(m, "Invalid NUMA node allocation handle success");
    } else {
        TEST_PASS(m, "Invalid NUMA node allocation should failed");
        passed++;
    }
    
    TEST_END(m, "NUMA locality tests: %d/%d passed", passed, total);
    return (passed == total) ? 0 : -1;
}

/**
 * test_numa_distance_aware_allocation - Tests allocation based on NUMA distance
 * 
 * This test verifies:
 * 1. Allocation preference for closer NUMA nodes
 * 2. Proper handling of NUMA distance matrix
 * 
 * Tests NUMA distance-aware allocation
 */
static int test_numa_distance_aware_allocation(struct seq_file *m)
{
    struct test_context *ctx = get_test_context();
    void *ptr;
    int num_nodes = num_online_nodes();
    int i;
    int passed = 0, total = 0;
    
    if (num_nodes < 2) {
        TEST_INFO(m, "NUMA distance test skipped: less than 2 nodes available");
        return 0;
    }
    
    // Limit to reasonable number
    if (num_nodes > 4) num_nodes = 4;
    
    // Test allocation on preferred node
    total++;
    ptr = l0_kmalloc(1024, ctx->test_node_id);
    if (ptr) {
        memset(ptr, ctx->test_node_id & 0xFF, 1024);
        l0_kfree(ptr);
        TEST_PASS(m, "Preferred NUMA node %d allocation successful", ctx->test_node_id);
        passed++;
    } else {
        TEST_FAIL(m, "Preferred NUMA node %d allocation failed", ctx->test_node_id);
    }
    
    // Test allocation on distant nodes
    for (i = 0; i < num_nodes; i++) {
        if (i != ctx->test_node_id) {
            total++;
            ptr = l0_kmalloc(1024, i);
            if (ptr) {
                memset(ptr, i & 0xFF, 1024);
                l0_kfree(ptr);
                TEST_PASS_DEBUG(m, "Distant NUMA node %d allocation successful", i);
                passed++;
            } else {
                TEST_FAIL(m, "Distant NUMA node %d allocation failed", i);
            }
        }
    }
    
    TEST_END(m, "NUMA distance-aware allocation tests: %d/%d passed", passed, total);
    return (passed == total) ? 0 : -1;
}

// Export test cases
static struct test_case test_numa_cases[] = {
    {"NUMA Node Allocation", test_numa_node_allocation},
    {"NUMA Node Migration", test_numa_node_migration},
    {"NUMA Locality", test_numa_locality},
    {"NUMA Distance Aware Allocation", test_numa_distance_aware_allocation},
};

DECLARE_TEST_MODULE(numa, test_numa_cases)